from langchain.chat_models import init_chat_model
from langchain_core.prompts import PromptTemplate
from pydantic import Field, BaseModel


class Date(BaseModel):
    year: int = Field(description="年")
    month: int = Field(description="月")
    day: int = Field(description="日")


json_schema = {
    "title": "Date",
    "description": "提取日期",
    "type": "object",
    "properties": {
        "year": {
            "type": "integer",
            "description": "year, YYYY",
        },
        "month": {
            "type": "integer",
            "description": "month, MM",
        },
        "day": {
            "type": "integer",
            "description": "day, DD",
        }
    }
}

# 今天是2025年10月8日,天气有雨

# 定义模型
llm = init_chat_model(
    model="deepseek-chat",
    model_provider="deepseek"
)

structured_output_llm = llm.with_structured_output(json_schema)

template = """
    提取用户输入文字中的日期:
    用户输入:{query},
"""
# 定义提示词
prompt = PromptTemplate(
    template=template,
    input_variables=["query"],

)

prompt_str = prompt.format(query="今天是2025年10月8日,天气有雨")
print(f"提示词:{prompt_str}")

ret = structured_output_llm.invoke(prompt_str)
print(ret)
print(ret.get("year"))
print(ret.get("month"))
print(type(ret))
